from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, ConfigurableFieldSpec
from langchain_openai import ChatOpenAI


chat_model = ChatOpenAI(
    openai_api_key="key",
    openai_api_base="https://api.moonshot.cn/v1",
    model="moonshot-v1-8k",
    temperature=0,
    request_timeout=60,
    max_retries=3,
)
prompt=ChatPromptTemplate.from_messages(
    [
        ("system","你是一个擅长{ability}的助手。"),
        MessagesPlaceholder(variable_name="history"),
        ("human","{input}"),
    ]
)
runnable=prompt|chat_model

store={}


def get_session_history(session_id:str)->BaseChatMessageHistory:
    if session_id not in store:
        store[session_id]=InMemoryChatMessageHistory()
    return store[session_id]

store_d={}
def get_session_history_with_conf(user_id:str,conversation_id:str)->BaseChatMessageHistory:
    if (user_id,conversation_id) not in store_d:
        store_d[(user_id,conversation_id)]=InMemoryChatMessageHistory()
    return store_d[(user_id,conversation_id)]

with_message_history = RunnableWithMessageHistory(
    runnable,get_session_history,input_messages_key="input",history_messages_key="history",
)
with_message_history_with_conf = RunnableWithMessageHistory(
    runnable,
    get_session_history,
    input_messages_key="input",
    history_messages_key="history",
    history_factory_config=[
        ConfigurableFieldSpec(
            id="user_id",
            annotation=str,
            name="用户 ID",
            description="用户的唯一标识符。",
            default="",
            is_shared=True,
        ),
        ConfigurableFieldSpec(
            id="conversation_id",
            annotation=str,
            name="对话 ID",
            description="对话的唯一标识符。",
            default="",
            is_shared=True,
        ),
    ],
)
res_conf=with_message_history.invoke(
    {"ability": "数学", "input": "你好"},
    config={"configurable": {"user_id": "123", "conversation_id": "1"}},
)
print(res_conf)

res=with_message_history.invoke(
    {"ability": "数学", "input": "余弦函数是什么意思？"},
    config={"configurable": {"session_id": "abc123"}},
)
print(res)

# 记住
res1=with_message_history.invoke(
    {"ability": "数学", "input": "什么？"},
    config={"configurable": {"session_id": "abc123"}},
)
print(res1)

# 新的 session_id --> 不记住。
res2=with_message_history.invoke(
    {"ability": "数学", "input": "什么？"},
    config={"configurable": {"session_id": "def234"}},
)
print(res2)
